# -*- coding: utf-8 -*-
import math

import pyrealsense2 as rs
import numpy as np
import cv2


class D435i:
    def __init__(self, ):
        """
        设置
        """
        self.w = 640
        self.h = 480
        self.fps = 30

        self.pipeline = rs.pipeline()  # 定义流程pipeline，创建一个管道
        config = rs.config()  # 定义配置config
        config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, self.fps)  # 配置depth流
        config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, self.fps)  # 配置color流

        # config.enable_stream(rs.stream.depth,  848, 480, rs.format.z16, 90)
        # config.enable_stream(rs.stream.color, 848, 480, rs.format.bgr8, 30)

        # config.enable_stream(rs.stream.depth,  1280, 720, rs.format.z16, 30)
        # config.enable_stream(rs.stream.color, 1280, 720, rs.format.bgr8, 30)

        pipe_profile = self.pipeline.start(config)  # streaming流开始

        # 创建对齐对象与color流对齐
        align_to = rs.stream.color  # align_to 是计划对齐深度帧的流类型
        self.align = rs.align(align_to)  # rs.align 执行深度帧与其他帧的对齐

    def get_aligned_images(self):
        """
        获取对齐图像帧与相机参数
        """
        frames = self.pipeline.wait_for_frames()  # 等待获取图像帧，获取颜色和深度的框架集
        aligned_frames = self.align.process(frames)  # 获取对齐帧，将深度框与颜色框对齐

        aligned_depth_frame = aligned_frames.get_depth_frame()  # 获取对齐帧中的的depth帧
        aligned_color_frame = aligned_frames.get_color_frame()  # 获取对齐帧中的的color帧

        #### 获取相机参数 ####
        depth_intrin = aligned_depth_frame.profile.as_video_stream_profile().intrinsics  # 获取深度参数（像素坐标系转相机坐标系会用到）
        color_intrin = aligned_color_frame.profile.as_video_stream_profile().intrinsics  # 获取相机内参

        #### 将images转为numpy arrays ####
        img_color = np.asanyarray(aligned_color_frame.get_data())  # RGB图
        img_depth = np.asanyarray(aligned_depth_frame.get_data())  # 深度图（默认16位）

        return color_intrin, depth_intrin, img_color, img_depth, aligned_depth_frame

    def get_3d_camera_coordinate(self, depth_pixel, aligned_depth_frame, depth_intrin):
        """
        获取随机点三维坐标
        """
        x = depth_pixel[0]
        y = depth_pixel[1]
        dis = aligned_depth_frame.get_distance(x, y)  # 获取该像素点对应的深度
        # print ('depth: ',dis)       # 深度单位是m
        camera_coordinate = rs.rs2_deproject_pixel_to_point(depth_intrin, depth_pixel, dis)
        # print ('camera_coordinate: ',camera_coordinate)
        return dis, camera_coordinate


def angle(camera_coordinate):
    pitch = camera_coordinate[0] / camera_coordinate[2]
    yaw = camera_coordinate[1] / camera_coordinate[2]
    return pitch, yaw


if __name__ == "__main__":
    d435i = D435i()
    while True:
        ''' 
        获取对齐图像帧与相机参数
        '''

        color_intrin, depth_intrin, img_color, img_depth, aligned_depth_frame = d435i.get_aligned_images()  # 获取对齐图像与相机参数

        ''' 
        获取随机点三维坐标
        '''
        depth_pixel = [320, 240]  # 设置随机点，以相机中心点为例
        dis, camera_coordinate = d435i.get_3d_camera_coordinate(depth_pixel, aligned_depth_frame, depth_intrin)
        print('depth: ', dis)  # 深度单位是mm
        print('camera_coordinate: ', camera_coordinate)

        '''
        获取旋转角度
        '''
        pitch, yaw = angle(camera_coordinate)
        print('pitch:', pitch, '--yaw:', yaw)

        ''' 
        显示图像与标注
        '''
        #### 在图中标记随机点及其坐标 ####
        cv2.circle(img_color, (320, 240), 8, [255, 0, 255], thickness=-1)
        cv2.putText(img_color, "Dis:" + str(dis) + " m", (40, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.2, [0, 0, 255])
        cv2.putText(img_color, "X:" + str(camera_coordinate[0]) + " m", (80, 80), cv2.FONT_HERSHEY_SIMPLEX, 1.2,
                    [255, 0, 0])
        cv2.putText(img_color, "Y:" + str(camera_coordinate[1]) + " m", (80, 120), cv2.FONT_HERSHEY_SIMPLEX, 1.2,
                    [255, 0, 0])
        cv2.putText(img_color, "Z:" + str(camera_coordinate[2]) + " m", (80, 160), cv2.FONT_HERSHEY_SIMPLEX, 1.2,
                    [255, 0, 0])

        #### 显示画面 ####
        cv2.imshow('RealSence', img_color)
        key = cv2.waitKey(1)
